
#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <string.h>
#include "utils.h"

#define PEPTIDESIZE 9

/* Argument Variables */
FILENAME	p_db;
FILENAME p_blm;
FILENAME p_fn;
WORD	p_alphabet;
int     p_seed;
int		n_db;

PARAM   param[] =
{
	"-db", VFNAME p_db, "File containing peptide database", "/home/projects/projects/vaccine/data/randompeps/randompeps_9mer_1000000.pep",	// Change this
    "-s", VINT      p_seed, "Seed [-1] Default, [0] Time [>0] Specific seed", "-1",
    "-dbn", VINT      n_db, "Number of database peptides to include", "10000",
    "-alp", VWORD	p_alphabet, "Amino acid alphabet", "ARNDCQEGHILKMFPSTWYV",
    "-blm", VFNAME p_blm, "Normalized blossum frequency matrix", "/home/projects/plichta/Jesper/aib-nn/data/blossum/blosum62.freq_rownorm",
    "-o", VFNAME p_fn, "Peptide list output file", "rand_data.pep",
	0
};

int compare (const void * a, const void * b)
{
	/* Used by qsort */
	return ( *(int*)a - *(int*)b );
}

void transform_blm(float **blm){
	/* Print blm */
	int i,j,k;

	/* Set diagonal to zero to force mutations */
	for (i=0; i<20; i++){
		blm[i][i] = 0.0;
	}

	/* Normalize blm */
	float sum, scl;
	for (i=0; i<20; i++){
		sum = 0.0;
		for (j=0; j<20; j++){
			sum += blm[i][j];
		}
		scl = 1.0 / sum;
		for (j=0; j<20; j++){
			blm[i][j] *= scl;
		}
	}

	/* Transform blm to accumulated probability */
	for (i=0; i<20; i++){
		for (j=1; j<19; j++){
			blm[i][j] += blm[i][j-1];
		}
		blm[i][19] = 1.0;
	}
}

PEPLIST *mutate_pep (PEPLIST *peplist, float **blm) {
	int first, middle, last, c_mut, pos, c;
	PEPLIST *pl;
	float p;

	for (pl=peplist; pl; pl=pl->next){
	//pl=peplist;
		pos = rand()%PEPTIDESIZE;
		c = strpos( p_alphabet, pl->pep[pos] );
		p = drand48();
		first = 0;
		last = 20-1;
		middle = (first+last)/2;

		//printf("pos=%d, first=%d, last=%d, middle=%d, c=%d, p=%3.2f\n", pos, first,last,middle,c,p);
		while ( first <= last ){
			if (p > blm[c][middle])
				first = middle + 1;
			else if (p == blm[c][middle]) {
				c_mut = middle;
				break;
			} else {
				/* p is smaller than middle */
				if (p > blm[c][middle-1]){
					/* p is in between middle and middle-1 */
					c_mut = middle;
					break;
				}
				last = middle-1;
			}
			middle = (first+last)/2;
		}
		if (first > last){
			printf("p=%3.2f not found in blm.\n", p);
			exit (1);
		}
		//printf("c=%d, c_mut=%d, c_mut=%c\n", c, c_mut, p_alphabet[c_mut]);
		pl->pep[pos] = p_alphabet[c_mut];
	}
	return peplist;
}

main(int argc, char *argv[])

{
	PEPLIST *db_peplist, *pl, **db_pep, *tra_peplist, *tst_peplist;
	float 	**blm;
	char	*bl_alphabet;
	int i,j,k;
	//"ARNDCQEGHILKMFPSTWYV";

	pparse( &argc, &argv, param, 2, "trainingfile" "testfile");

	db_peplist = peplist_read( p_db );
	tra_peplist = peplist_read( argv[1] );
	tst_peplist = peplist_read( argv[2] );
	blm = read_realblosum( p_blm, &bl_alphabet );

	if ( strcmp( bl_alphabet, p_alphabet ) ) {
		printf( "Error. Bl alphabet %s not equal to p_alphabet %s\n",
			bl_alphabet, p_alphabet );
		exit( 1 );
	}

	transform_blm(blm);

	srand48(time(NULL));
	srand(time(NULL));

	PEPLIST *muttra_peplist, *muttst_peplist;

	muttra_peplist = mutate_pep(tra_peplist, blm);
	muttst_peplist = mutate_pep(tst_peplist, blm);

	if ( p_seed >= 0 )
		setseed( p_seed );		// not used

    if ( db_peplist == NULL ) {
            printf( "No Peptides sequences read from file %s\n", argv[1] );
            exit( 1 );
    }
    int db_size, *sel;

	/* Count size of peptide database */    
    for (pl=db_peplist, db_size=0; pl; pl=pl->next, db_size++){}

    /* Determine which peptides from database to include */
    sel = ivector(0, n_db-1);
	
    for (i=0; i<n_db; i++){
    	sel[i] = rand()%(db_size-1);	// Fix seed used here.
    }

    /* Sort database peptides */
    qsort(sel, n_db, sizeof(int), compare);

    /* Check for duplicates */
    for (i=1; i<n_db; i++){
    	if(sel[i] == sel[i-1]){
    		for (k=3; k>0; k--){
    			if (k==db_size){
    				sel[i+k] = rand()%(db_size-1);	// Pick new
    			}
    			if (sel[i] == sel[i+k])
    				sel[i+k] = sel[i+k] + k;
    		}
    		sel[i]++;
    	}
    }

	/* Alloc db_pep list */
    db_pep = malloc(n_db * sizeof ( PEPLIST* ) );
    for (i=0; i<n_db; i++)
    	db_pep[i] = NULL;

    /* Make list of pointers to peptides that should be included */
    int db_count = 0;
    for (pl=db_peplist, i=0; pl; pl=pl->next, i++){
    	if (i != sel[db_count]){
    		continue;
    	} else {
    		db_pep[db_count] = pl; 
    		db_count++;
    		if(db_count == n_db)	// The number of database peptides to include has been reached.
    			break;				
    	}
    }


    /* Print included peptides */

    PEPLIST *p1, *p2, *p3, *p4;
    int pep_count;
    for (pep_count=0, p1=tra_peplist; p1; p1=p1->next, pep_count++){}
	for (p2=tst_peplist; p2; p2=p2->next, pep_count++){}
	for (p3=muttra_peplist; p3; p3=p3->next, pep_count++){}
	for (p4=muttst_peplist; p4; p4=p4->next, pep_count++){}
	pep_count += n_db;
	
	printf("# Number of peptides generated: %d\n", pep_count);

	FILE 	*fp;
	int		fc;

	if ( ( fp = stream_output( p_fn, &fc, 0)) == NULL) {
		printf( " Error. Can't open %s\n", p_fn);
		exit(1);
	}

	fp = fopen(p_fn, "a");
	
    for (p1=tra_peplist; p1; p1=p1->next)
    	fprintf(fp, "%s\n", p1->pep);
	for (p2=tst_peplist; p2; p2=p2->next, pep_count++)
		fprintf(fp, "%s\n", p2->pep);
	for (p3=muttra_peplist; p3; p3=p3->next, pep_count++)
		fprintf(fp, "%s\n", p3->pep);
	for (p4=muttst_peplist; p4; p4=p4->next, pep_count++)
		fprintf(fp, "%s\n", p4->pep);
	for (i=0; i<n_db; i++)
    	fprintf(fp, "%s\n", db_pep[i]->pep); 

	exit( 0 );
}